///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file psi3_interface.cpp
///
/// \brief   contains definitions of PSI3 access functions
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/blas_lapack.h>
#include <src/psi3_interface.h>

#ifdef HAVE_PSI3

// include headers needed for PSI3
extern "C" {
#include <stdio.h>
#include <stdlib.h>
#include <libciomr/libciomr.h>
#include <psifiles.h>
#include <libchkpt/chkpt.h>
#include <libpsio/psio.h>
#include <libiwl/iwl.h>
#include <libipv1/ip_lib.h>
#include <libqt/qt.h>
#include <math.h>
}

// define global variables for PSI3 libraries
extern "C" {
  char *psi_file_prefix;
  FILE *infile, *outfile;
  char *gprgid() { char *prgid = ":PSICHEM"; return(prgid); }
}

#endif

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   reads the one and two electron integrals from PSI3
///
/// \param[in]   userinp  user input options
/// \param[out]  enuc     the nuclear repulsion energy
/// \param[out]  oei      the one electron integrals, size = na*na
/// \param[out]  tei      the two electron integrals, size = na*(na+1) * ( na*(na+1) + 2 ) / 8
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void pcps::get_psi3_ints(const pcps::Input & userinp, double & enuc, double * const oei, double * const tei) {

  #ifndef HAVE_PSI3
  throw pcps::Exception("cannot read psi3 integrals.  please recompile with psi3.");
  #endif

  #ifdef HAVE_PSI3

  // initialize psi3
  const int argc = 0;
  char ** argv;
  //const int errcod = psi_start(argc, argv, 0);
  const int errcod = psi_start(&infile, &outfile, &psi_file_prefix, argc, argv, 0);
  if ( errcod != PSI_RETURN_SUCCESS )
    throw pcps::Exception( (boost::format("psi_start failed with code %i in pcps::get_psi3_ints") % errcod).str() );
  //ffile(&infile, "psichem.inp", 2);
  //ffile(&outfile, "psichem.out", 1);
  ip_cwk_add(gprgid());
  tstart(outfile);
  psio_init();
  psio_ipv1_config();

  // open the chkpt file
  chkpt_init(PSIO_OPEN_OLD);

  // load the nuclear repulsion energy
  enuc = chkpt_rd_enuc();

  // get numbers of orbitals
  const int no = chkpt_rd_nso();       // total
  const int nc = userinp.n_cor_orb();  // core
  const int na = userinp.n_act_orb();  // active

  // check that the numbers of orbitals are sane
  if ( nc + na > no )
    throw pcps::Exception("too many core and/or active orbitals.  nc = %i, na = %i, no = %i") % nc % na % no;

  // check that the number of active orbitals matches the number of lattice sites
  if ( userinp.nsites() != 2*na )
    throw pcps::Exception("number of active orbitals must equal lattice sites:  na = %i, nsites = %i") % na % (userinp.nsites()/2);

  // close the chkpt files
  chkpt_close();

  // create temporary vectors into which we will load the integrals
  std::vector<double> temp_int1(no*no, 0.0);
  std::vector<double> temp_int2( no*(no+1)/2 * ( no*(no+1)/2 + 1 ) / 2, 0.0);

  std::vector<double> ao_overlap(no*no, 0.0);

  // get the number of elements in the one electron integral's lower triangle
  const int ntri = no*(no+1)/2;

  // create a vector used in loading one electron data from PSI3
  std::vector<double> lower_triangle(ntri, 0.0);

  // load the one electron integrals in lower triangle format
  if (userinp.s_half_orth()) {
    lower_triangle.assign(lower_triangle.size(), 0.0);
    iwl_rdone(PSIF_OEI, PSIF_SO_T, &lower_triangle.at(0), ntri, 0, 0, outfile);
    pcps::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));
    lower_triangle.assign(lower_triangle.size(), 0.0);
    iwl_rdone(PSIF_OEI, PSIF_SO_V, &lower_triangle.at(0), ntri, 0, 0, outfile);
    pcps::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));
    lower_triangle.assign(lower_triangle.size(), 0.0);
    iwl_rdone(PSIF_OEI, PSIF_SO_S, &lower_triangle.at(0), ntri, 0, 0, outfile);
    pcps::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &ao_overlap.at(0));
  } else {
    iwl_rdone(PSIF_OEI, PSIF_MO_OEI, &lower_triangle.at(0), ntri, 0, 0, outfile);
    pcps::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));
  }

  // precompute offset array used in loading two electron integrals
  std::vector<int> ioff(ntri*ntri, 0);
  for (int i = 0; i < ioff.size(); i++)
    ioff[i] = i*(i+1)/2;

  // load two electron integrals
  if (userinp.s_half_orth())
    iwl_rdtwo(PSIF_SO_TEI, &temp_int2.at(0), &ioff.at(0), no, 0, 0, 0, outfile);
  else
    iwl_rdtwo(PSIF_MO_TEI, &temp_int2.at(0), &ioff.at(0), no, 0, 0, 0, outfile);

  // finalize psi3
  psio_done();
  //psi_stop();
  psi_stop(infile, outfile, psi_file_prefix);

  // perform S^(-1/2) orthogonalization if requested
  if (userinp.s_half_orth())
    pcps::orthogonalize_orbs_by_overlap(no, ao_overlap, temp_int1, temp_int2);

  // extract active space one electron integrals
  for (int p = 0; p < na; p++)
  for (int q = 0; q < na; q++) {
    oei[na*p+q] = temp_int1[no*(p+nc)+(q+nc)];
    for (int i = 0; i < nc; i++)
      oei[na*p+q] +=   2.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, i, p+nc, q+nc)]
                     - 1.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, p+nc, i, q+nc)];
  }

  // extract active space two electron integrals
  for (int p = 0; p < na; p++)
  for (int q = 0; q <= p; q++)
  for (int r = 0; r < na; r++)
  for (int s = 0; s <= r; s++) {
    const int tmp_cmpd = pcps::Hamiltonian<double>::tei_compound_index(p+nc, q+nc, r+nc, s+nc);
    const int act_cmpd = pcps::Hamiltonian<double>::tei_compound_index(p, q, r, s);
    tei[act_cmpd] = temp_int2[tmp_cmpd];
  }

  // add core orbital contribution to the nuclear energy (nuclear energy is used to also hold core energy)
  for (int i = 0; i < nc; i++)
    enuc += 2.0 * temp_int1[no*i+i] + temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, i, i, i)];
  for (int i = 0; i < nc; i++)
  for (int j = 0; j <  i; j++)
    enuc +=   4.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, i, j, j)]
            - 2.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, j, i, j)];

  #endif

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   throws an error because PSI3 doesn't use complex integrals
///
/// \param[in]   userinp  user input options
/// \param[out]  enuc     the nuclear repulsion energy
/// \param[out]  oei      the one electron integrals, size = na*na
/// \param[out]  tei      the two electron integrals, size = na*(na+1) * ( na*(na+1) + 2 ) / 8
///
///////////////////////////////////////////////////////////////////////////////////////////////////


void pcps::get_psi3_ints(const pcps::Input & userinp, double & enuc, std::complex<double> * const oei,
                                                                     std::complex<double> * const tei) {

  throw pcps::Exception("Cannot read complex integrals from PSI3.  Please use scalar_type = real");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   converts from the AO 1- and 2-electron integrals to the S^(-1/2) orthogonalized basis
///          1- and 2- electron integrals
///
/// \param[in]      n        the number of ao orbitals
/// \param[in]      overlap  the ao overlap matrix, size n*n
/// \param[in,out]  int1     on entry, the ao one electron integrals, size n*n
///                          on exit, the one electron integrals in the new basis
/// \param[in,out]  int2     on entry, the ao two electron integrals,
///                          size  n*(n+1)/2 * ( n*(n+1)/2 + 1 ) / 2
///                          on exit, the two electron integrals in the new basis
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void pcps::orthogonalize_orbs_by_overlap(const int n,
                                         const std::vector<double> & overlap,
                                         std::vector<double> & int1,
                                         std::vector<double> & int2)
{

  std::vector<double> s_neg_half(overlap);
  std::vector<double> intermed_0(n*n, 0.0);
  std::vector<double> intermed_1(n*n, 0.0);

  std::vector<double> big2_0(n*n*n*n, 0.0);
  std::vector<double> big2_1(n*n*n*n, 0.0);

  // preparation for diagonalization of overlap matrix
  int info;
  const int lwork = 5*n;
  std::vector<double> work(lwork, 0.0);
  std::vector<double> rwork(3*n, 0.0);
  std::vector<double> eval(n, 0.0);

  // diagonalize the overlap matrix, after which the eigenvectors will be in the rows the matrix
  pcps::xsyev('V', 'U', n, &s_neg_half.at(0), n, &eval.at(0), &work.at(0), lwork, &rwork.at(0), info);
  if (info != 0)
    throw pcps::Exception("pcps::xsyev failed with error code %i in pcps:::orthogonalize_orbs_by_overlap") % info;

  // compute the S^(-1/2) matrix
  for (int i = 0; i < n; i++) {
    const double x = 1.0 / std::sqrt(eval[i]);
    for (int j = 0; j < n; j++) {
      intermed_0[i*n+j] = s_neg_half[j*n+i];
      intermed_1[i*n+j] = s_neg_half[i*n+j] * x;
    }
  }
  pcps::xgemm('N', 'N', n, n, n, 1.0, &intermed_1.at(0), n, &intermed_0.at(0), n, 0.0, &s_neg_half.at(0), n);

  // print the orthonormal orbital coefficients
  std::cout << boost::format("orthonormal basis orbitals coefficients for S^(-1/2) basis") << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      std::cout << boost::format(" %10.6f") % s_neg_half[i*n+j];
    std::cout << std::endl;
  }
  std::cout << std::endl;

  // convert the one electron integrals
  pcps::xgemm('N', 'N', n, n, n, 1.0, &s_neg_half.at(0), n, &int1.at(0), n, 0.0, &intermed_0.at(0), n);
  pcps::xgemm('N', 'N', n, n, n, 1.0, &intermed_0.at(0), n, &s_neg_half.at(0), n, 0.0, &int1.at(0), n);

  // print the one electron integrals
  std::cout << boost::format("one electron integrals for S^(-1/2) basis") << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      std::cout << boost::format(" %10.6f") % int1[i*n+j];
    std::cout << std::endl;
  }
  std::cout << std::endl;

  // convert the two electron integrals

  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
  for (int r = 0; r < n; r++)
  for (int s = 0; s < n; s++)
    big2_0[n*n*n*p+n*n*q+n*r+s] = int2[pcps::Hamiltonian<double>::tei_compound_index(p,q,r,s)];

  struct ConvertAndShift {
    static void foo(const int n, const std::vector<double> & s_neg_half, double * p0, double * p1) {
      pcps::xgemm('N', 'N', n*n*n, n, n, 1.0, p0, n*n*n, &s_neg_half.at(0), n, 0.0, p1, n*n*n);
      for (int p = 0; p < n; p++)
      for (int q = 0; q < n; q++)
      for (int r = 0; r < n; r++)
      for (int s = 0; s < n; s++)
        p0[n*n*n*p+n*n*q+n*r+s] = p1[n*n*n*q+n*n*r+n*s+p];
    }
  };
  ConvertAndShift::foo(n, s_neg_half, &big2_0.at(0), &big2_1.at(0));
  ConvertAndShift::foo(n, s_neg_half, &big2_0.at(0), &big2_1.at(0));
  ConvertAndShift::foo(n, s_neg_half, &big2_0.at(0), &big2_1.at(0));
  ConvertAndShift::foo(n, s_neg_half, &big2_0.at(0), &big2_1.at(0));

  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
  for (int r = 0; r < n; r++)
  for (int s = 0; s < n; s++)
    int2[pcps::Hamiltonian<double>::tei_compound_index(p,q,r,s)] = big2_0[n*n*n*p+n*n*q+n*r+s];

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   expands the psi3 lower triangle storage for one electron integrals
///
///   This function adds the one electron integral data stored in the lower triangle to the one
///   electron integral data stored in the full matrix.
///
/// \param[in]      n        the number of spatial orbitals
/// \param[in]      lt       the one electron integrals in lower triangle storage, size n*(n+1)/2
/// \param[in,out]  oei      the one electron integrals in full-matrix storage, size = n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void pcps::unpack_psi3_lower_triangle(const int n, const double * const lt, double * const oei) {

  for (int i = 0; i < n; i++) {

    // diagonal elements
    oei[i*n+i] += lt[ i*(i+1)/2 + i ];

    // off diagonal elements
    for (int j = 0; j < i; j++) {
      oei[i*n+j] += lt[ i*(i+1)/2 + j ];
      oei[j*n+i] += lt[ i*(i+1)/2 + j ];
    }

  }

}
